import argparse
import os
import sys
sys.path.append(os.getcwd())
sys.path.append("..")
import pandas as pd
from scipy.stats import spearmanr
from tqdm import tqdm
from src.utils import mutant_filter


def process(args):
    if args.trg_proteins is not None:
        proteins = sorted([p.strip() for p in open(args.trg_proteins, "r").readlines()])
    else:
        proteins = sorted(os.listdir(args.dataset_input))
    dataset_name = args.dataset_input.split("/")[-2]
    
    res = {"name": proteins, "count": []}
    models = [m.split("/")[-1] for m in args.model_location]
    for model in models:
        res[model] = []
    for protein in tqdm(proteins):
        for model in models:
            truth_score = pd.read_table(f"{args.dataset_input}/{protein}/experiments/{protein}.tsv")
            pred_score = pd.read_table(f"{args.dataset_input}/{protein}/predictions/{protein}.{model}.tsv")
            truth_score = mutant_filter(truth_score, args.mutant_site).dropna()
            pred_score = mutant_filter(pred_score, args.mutant_site).loc[truth_score.index]
            assert len(truth_score) == len(pred_score), f"{protein}: {len(truth_score)} {len(pred_score)}"
            score = spearmanr(truth_score["score"], pred_score["score"])[0]
            res[model].append(score)
        count = len(truth_score)
        res["count"].append(count)
    df = pd.DataFrame(res)
    
    for model in models:
        print(f"{model} average spearmanr: ", df[model].mean())
    if args.output_path is None:
        out_file = os.path.join(args.dataset_input, dataset_name+f"_{args.mutant_site}_spearmanr.csv")
    else:
        out_file = args.output_path
    df.to_csv(out_file, index=False)
    print(f"Saving to {out_file}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--trg_proteins", type=str, default=None, help="select data point in txt file")
    parser.add_argument("--dataset_input", type=str, default="/nvme/tyang/workspace/Proteus/data")
    parser.add_argument("--mutant_site", type=int, default=0)
    parser.add_argument("--model_location", type=str, nargs="+", default=["esm1v_avg"])
    parser.add_argument("--output_path", type=str, default=None)
    args = parser.parse_args()
    
    process(args)